# utils/visualization.py
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
from sklearn.manifold import TSNE


def visualize_prototypes(model, data_loader, device, num_classes=11):
    """可视化原型和特征分布"""
    model.eval()
    all_features = []
    all_labels = []

    with torch.no_grad():
        for X, y in data_loader:
            X, y = X.to(device), y.to(device)
            _, _, features = model(X)
            all_features.append(features.cpu().numpy())
            all_labels.append(y.cpu().numpy())

    # 获取原型
    prototypes = model.class_prototypes.cpu().numpy()

    # 合并特征
    features = np.concatenate(all_features)
    labels = np.concatenate(all_labels)

    # 添加原型
    features = np.vstack([features, prototypes])
    proto_labels = np.arange(num_classes) + num_classes
    labels = np.concatenate([labels, proto_labels])

    # t-SNE降维
    tsne = TSNE(n_components=2, random_state=42)
    embeddings = tsne.fit_transform(features)

    # 绘制
    plt.figure(figsize=(12, 10))

    # 绘制样本点
    scatter = plt.scatter(
        embeddings[:-num_classes, 0],
        embeddings[:-num_classes, 1],
        c=labels[:-num_classes],
        cmap='tab20',
        alpha=0.6,
        label='Samples'
    )

    # 绘制原型点
    proto_scatter = plt.scatter(
        embeddings[-num_classes:, 0],
        embeddings[-num_classes:, 1],
        c=labels[-num_classes:] - num_classes,
        cmap='tab20',
        marker='*',
        s=300,
        edgecolors='black',
        linewidths=1,
        label='Prototypes'
    )

    plt.title('t-SNE Visualization of Features and Prototypes')
    plt.colorbar(scatter, label='Class')
    plt.legend()
    plt.tight_layout()
    plt.savefig('prototype_visualization.png')
    plt.close()